import matplotlib.pyplot as plt
import mpl_finance as mpf
import pandas as pd
import numpy as np

import shujushouji.base  # 股票数据文件解析
import shujujisuan.mathutils  # 计算涨停价、跌停价、均价

ax_label_size = 30  # x轴标签字体大小
ay_label_size = 45  # y轴标签字体大小


def paint_rik(rik_paras):

    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

    code = rik_paras['code']  # 股票代码
    name = rik_paras['name']  # 股票名称
    curr_date = rik_paras['date']  # 股票异动日期
    onmarket_date = rik_paras['onmarket_date']  # 股票上市日期
    begin_date = rik_paras['begin_date']  # 日K线开始日期
    end_date = rik_paras['end_date']  # 日K线结束日期
    all_date_list = rik_paras['all_date_list']  # 股市所有交易日
    file_name = rik_paras['file_name']  # 日K图保存文件名

    fig = plt.figure(  # 添加fig对象
        figsize=(60, 80),  # 设置fig大小，长和宽，单位为英寸
        dpi=50,  # 每英寸的像素点数
        facecolor="white")
    gs = fig.add_gridspec(
        9, 1,  # 9行1列的网格
        left=0.01, right=0.96, bottom=0.03, top=0.99,  # 设置页边距
        wspace=0.02, hspace=0.12,  # 设置行与行之间的距离，列与列之间的距离
        height_ratios=[0.3, 0.18, 0.45, 0.15, 0.45, 0.15, 0.15, 0.55, 0.23])  # 这边是9行，设置每行占的高度的比例

    # 0、股票的入选原因及基础信息
    ax_xinxi = fig.add_subplot(gs[0, 0])
    ax_xinxi.text(0.01, 0.5, '0、股票的入选原因及基础信息',
                  ha="left", va='center', fontsize='100')

    # 1、股票的业绩数据，以表格形式呈现
    ax_yeji = fig.add_subplot(gs[1, 0])
    ax_yeji.text(0.01, 0.5, '1、股票的业绩数据，以表格形式呈现',
                 ha="left", va='center', fontsize='100')

    # 2、指数的日K图（包括20日、10日、5日均线）
    ax_zhishu_rik = fig.add_subplot(gs[2, 0])
    ax_zhishu_rik.text(0.01, 0.5, '2、指数的日K图（包括20日、10日、5日均线）',
                       ha="left", va='center', fontsize='100')

    # 3、指数的成交额柱状图
    ax_zhishu_cje = fig.add_subplot(gs[3, 0])
    ax_zhishu_cje.text(0.01, 0.5, '3、指数的成交额柱状图',
                       ha="left", va='center', fontsize='100')

    stock_data_dic = calc_dataframe_new(  # 获取股票行情数据用于绘图
        code, name, curr_date, onmarket_date, begin_date, end_date, all_date_list)
    stock_data = pd.DataFrame(stock_data_dic)  # 将股票行情数据转换成DataFrame格式

    # 4、个股的日K图（包括20日、10日、5日均线）
    ax_rik = fig.add_subplot(gs[4, 0])
    gegu_rik(ax_rik, stock_data)

    # 5、个股的换手率柱状图（柱状图上的标签为当日的成交额）
    ax_huanshoulv = fig.add_subplot(gs[5, 0])  # 第5行0列的坐标轴
    ax_huanshoulv.text(0.01, 0.5, '5、个股的换手率柱状图（柱状图上的标签为当日的成交额）',
                       ha="left", va='center', fontsize='100')

    # 6、个股的大资金折线图（红色折线图为大资金流入，绿色折线图为大资金流出）
    ax_dzj = fig.add_subplot(gs[6, 0])
    ax_dzj.text(0.01, 0.5, '6、个股的大资金折线图（红色折线图为大资金流入，绿色折线图为大资金流出）',
                ha="left", va='center', fontsize='100')

    # 7、个股的十大流通股东数据，，以表格形式呈现
    ax_gudong = fig.add_subplot(gs[7, 0])
    ax_gudong.text(0.01, 0.5, '7、个股的十大流通股东数据，以表格形式呈现',
                   ha="left", va='center', fontsize='100')

    # 8、留白，待后续补充
    ax_liubai = fig.add_subplot(gs[8, 0])
    ax_liubai.text(0.01, 0.5, '8、留白，待后续补充',
                   ha="left", va='center', fontsize='100')

    plt.savefig(file_name)  # 保存图片
    plt.close(fig)


def gegu_rik(ax_rik, stock_data):
    mpf.candlestick2_ochl(  # 调用candlestick2_ochl画日K图
        ax_rik,  # 在这套坐标轴内画日K图
        stock_data.open, stock_data.close,  # 开盘价和收盘价
        stock_data.high, stock_data.low,  # 最高价和最低价
        width=0.8, colorup='red', colordown='green', alpha=1.0)  # 收盘价大于开盘价则红柱，收盘价小于开盘价则绿柱
    ax_rik.grid(axis='x', linestyle='-.')  # 添加x轴方向的网格
    ax_rik.grid(axis='y', linestyle='-.')  # 添加y轴方向的网格
    ax_rik.yaxis.tick_right()  # 将y轴的数值显示在右边
    ax_rik.set_xlim(0, len(stock_data), 1)  # 设置x轴的范围
    ax_rik.set_xticks(range(0, len(stock_data), 1))  # 设置x轴的刻度
    ax_rik.axes.set_xticklabels(stock_data.datelabel)

    for label in ax_rik.xaxis.get_ticklabels():  # 设置x轴标签的字体大小
        label.set_fontsize(ax_label_size)
    for label in ax_rik.yaxis.get_ticklabels():  # 设置y轴标签的字体大小
        label.set_fontsize(ay_label_size)

    # 绘制移动平均线图，分别是5日平均线、10日平均线、20日平均线
    ax_rik.plot(
        np.arange(0, len(stock_data)),  # [x]的数据范围
        stock_data.jun5,  # [y]的数据范围
        'black',  # 均线颜色
        label='M5',  # 均线标签
        lw=2.5)  # 均线宽度
    ax_rik.plot(np.arange(0, len(stock_data)), stock_data.jun10, 'orange', label='M10', lw=2.5)
    ax_rik.plot(np.arange(0, len(stock_data)), stock_data.jun20, 'purple', label='M20', lw=2.5)


def calc_dataframe_new(
        code,  # 股票代码
        name,  # 股票名称
        curr_date,  # 股票异动日期
        onmarket_date,  # 股票的上市日期
        begin_date,  # 绘制的日K线起始日期
        end_date,  # 绘制的日K线结束日期
        all_date_list):  # 所有交易日期

    data_frame = {
        'open': [], 'close': [], 'high': [], 'low': [],  # 蜡烛图需要的四个价格
        'huanshoulv': [], 'chengjiaoe': [],  # 换手率柱状图需要
        'jingbuy': [], 'jingsell': [],  # 大资金折线图需要
        'datelabel': [],  # x轴的日期标签
        'cjelabel': [],  # 用在大资金折线图，但感觉没用好
        'jun5': [], 'jun10': [], 'jun20': []  # 三条均线的每日数据，
    }

    onmarket_date_index = 0
    if onmarket_date in all_date_list:
        onmarket_date_index = all_date_list.index(onmarket_date)

    begin_date_index = all_date_list.index(begin_date)
    end_date_index = all_date_list.index(end_date)

    stock_data_all_dic, stock_data_all_list = shujushouji.base.get_stock_data_by_code(code)  # 获取股票历史行情数据
    last_date_stock_data_dic = {}
    if onmarket_date in stock_data_all_dic.keys():
        last_date_stock_data_dic = stock_data_all_dic[onmarket_date]

    jun5 = 0.0  # 5日均价
    jun10 = 0.0  # 10日均价
    jun20 = 0.0  # 20日均价
    if begin_date_index > onmarket_date_index:  # 如果日k线统计日期晚于股票上市日期，初始化last_date_stock_data_dic, jun5, jun10, jun20
        for index in range(begin_date_index - 1, onmarket_date_index, -1):
            index_date = all_date_list[index]
            if index_date in stock_data_all_dic.keys():
                last_date_stock_data_dic = stock_data_all_dic[index_date]
                jun5 = shujushouji.base.calc_avg_price(5, index_date, onmarket_date, all_date_list,
                                                       stock_data_all_dic)
                jun10 = shujushouji.base.calc_avg_price(10, index_date, onmarket_date, all_date_list,
                                                        stock_data_all_dic)
                jun20 = shujushouji.base.calc_avg_price(20, index_date, onmarket_date, all_date_list,
                                                        stock_data_all_dic)
                break

    for index in range(begin_date_index, end_date_index + 1):
        index_date = all_date_list[index]
        if index_date in stock_data_all_dic.keys():  # 当前交易日存在交易行情数据
            data_dic = stock_data_all_dic[index_date]
            data_frame['open'].append(data_dic['open'])
            data_frame['close'].append(data_dic['close'])
            data_frame['high'].append(data_dic['high'])
            data_frame['low'].append(data_dic['low'])
            data_frame['huanshoulv'].append(data_dic['huanshoulv'])
            data_frame['chengjiaoe'].append(data_dic['chengjiaoe'])
            data_frame['jingbuy'].append(data_dic['sbigin'] + data_dic['bigin'])
            data_frame['jingsell'].append(- data_dic['sbigout'] - data_dic['bigout'])

            # 获取当天的涨停价和跌停价
            zt_price, dt_price = shujujisuan.mathutils.get_zt_dt(code, name, data_dic['lastclose'])

            # 给日期标签加上一些记号，便于识别，异动当天日期后缀为B(begin)，涨停的交易日标签后缀为R，跌停的交易日标签后缀为V
            if index_date == curr_date:
                data_frame['datelabel'].append(index_date[-1] + 'B')
            else:
                if data_dic['close'] == zt_price:
                    data_frame['datelabel'].append(index_date[-1] + 'R')
                elif data_dic['close'] == dt_price:
                    data_frame['datelabel'].append(index_date[-1] + 'V')
                else:
                    data_frame['datelabel'].append(index_date[-2:])

            data_frame['cjelabel'].append(int(data_dic['chengjiaoe'] + 0.5))
            jun5 = shujushouji.base.calc_avg_price(5, index_date, onmarket_date, all_date_list, stock_data_all_dic)
            jun10 = shujushouji.base.calc_avg_price(10, index_date, onmarket_date, all_date_list, stock_data_all_dic)
            jun20 = shujushouji.base.calc_avg_price(20, index_date, onmarket_date, all_date_list, stock_data_all_dic)
            data_frame['jun5'].append(jun5)
            data_frame['jun10'].append(jun10)
            data_frame['jun20'].append(jun20)

            last_date_stock_data_dic = data_dic
        else:  # 如果这个日期股票还未上市的话就用上市当天的数据;如果这个交易日停牌的话，就用停牌前一天的数据
            data_frame['open'].append(last_date_stock_data_dic['open'])
            data_frame['close'].append(last_date_stock_data_dic['open'])
            data_frame['high'].append(last_date_stock_data_dic['open'])
            data_frame['low'].append(last_date_stock_data_dic['open'])
            data_frame['huanshoulv'].append(0)
            data_frame['chengjiaoe'].append(0)
            data_frame['jingbuy'].append(0)
            data_frame['jingsell'].append(0)
            data_frame['datelabel'].append(index_date[-2:])
            data_frame['cjelabel'].append(0)
            data_frame['jun5'].append(jun5)
            data_frame['jun10'].append(jun10)
            data_frame['jun20'].append(jun20)
    return data_frame


def main():
    all_date_list = shujushouji.base.get_all_date()  # 获取所有交易日期
    latest_date_index = len(all_date_list) - 1  # 最新交易日的序号
    curr_date = '2022-11-11'
    curr_date_index = all_date_list.index(curr_date)
    rik_days = 120  # 日k线绘制的交易日数
    after_days = 60  # 最多绘制异动日期后多少个交易日的K线

    # 根据股票异动日期, rik_days, after_days 初始化日K图的开始绘制日期和结束绘制日期
    end_date_index = latest_date_index  # 先确定k线结束日期
    if (latest_date_index - curr_date_index) >= after_days:
        end_date_index = curr_date_index + after_days
    begin_date_index = end_date_index - rik_days

    code = '600519'
    code_name_dic, code_onmarket_dic = shujushouji.base.get_all_code_and_onmarket_date()
    onmarket_date = all_date_list[0]
    if code in code_onmarket_dic.keys():
        onmarket_date = code_onmarket_dic[code]

    rik_paras = {'code': code,  # 股票代码
                 'name': code_name_dic[code],  # 股票名称
                 'date': curr_date,  # 股票异动日期
                 'onmarket_date': onmarket_date,  # 股票上市日期
                 'begin_date': all_date_list[begin_date_index],  # 日K线开始日期
                 'end_date': all_date_list[end_date_index],  # 日K线结束日期
                 'all_date_list': all_date_list,  # 股市所有交易日
                 'file_name': 'D:\\stock\\fenxi\\' + code + '-' + code_name_dic[code] + '-' + curr_date + '.jpg'  # 日K图保存文件名
                 }
    paint_rik(rik_paras)


if __name__ == '__main__':
    main()
